from typing import Callable

import numpy as np
from scipy.spatial.transform import Rotation as R
from sklearn.metrics import pairwise_distances

import carla
from dora import DoraStatus
from dora_utils import LABELS, Decision, Action, Status, wrap_angle
from frenet_optimal_trajectory_planner.FrenetOptimalTrajectory import (
    fot_wrapper,
)
import matplotlib.pyplot as plt

# Local Planner
LOCAL_PLANNER_NUM_GLOBAL_WAYPOINTS_AHEAD = 20
LOCAL_PLANNER_MAX_SPEED = 10  # m/s
LOCAL_PLANNER_MAX_ACCEL = 8  # m/s^2
LOCAL_PLANNER_USE_LANE_OBSTACLES = True

# Obstacles
OBSTACLE_DIST_THRESHOLD = 2.5  # m
OBSTACLE_ANGLE_THRESHOLD = np.pi / 2  # rad
OBSTACLE_CONFIDENCE_THRESHOLD = 50
OBSTACLE_RADIUS = 0.5  # m
OBSTACLE_RADIUS_TANGENT = 0.5  # m
OBSTACLE_LANE_RADIUS = 0.15  # m
OBSTACLE_IGNORE_LABELS = [
    "fire hydrant",
    "bench",
    "traffic light",
]
OBSTACLE_NONCROSS_LANE_TYPE = [
    carla.LaneMarkingType.NONE,
    carla.LaneMarkingType.Solid,
    carla.LaneMarkingType.SolidSolid,
]

# Vehicle
VEHICLE_DT = 1.0 / 20  # s

# Debug
DEBUG = False

if DEBUG:
    import pickle, os

    filename = "/home/dora/workspace/simulate/fot_debug.pkl"
    if os.path.exists(filename):
        os.remove(filename)


def nearest_k_points(points, position, k):
    """Get nearest k points"""
    distances = np.linalg.norm(points - position, axis=1)
    nearest_idx = np.argpartition(distances, k)[:k]
    nearest_k_pts = points[nearest_idx]

    return nearest_k_pts


def get_noncross_boundary(lane_waypoints):
    """Get non-cross boundary list"""
    noncross_boundary = []

    for waypoint in lane_waypoints:
        left_lane_marking = waypoint.left_lane_marking
        right_lane_marking = waypoint.right_lane_marking
        noncross_left = False
        noncross_right = False
        if left_lane_marking.type in OBSTACLE_NONCROSS_LANE_TYPE:
            noncross_left = True
        if right_lane_marking.type in OBSTACLE_NONCROSS_LANE_TYPE:
            noncross_right = True

        # use the yaw (orientation) to compute the tangent direction
        yaw_rad = np.deg2rad(waypoint.transform.rotation.yaw)
        dx = np.cos(yaw_rad)
        dy = np.sin(yaw_rad)

        # adjust for pitch and roll (optional, the effect is small)
        pitch_rad = np.deg2rad(waypoint.transform.rotation.pitch)
        roll_rad = np.deg2rad(waypoint.transform.rotation.roll)
        dx = dx * np.cos(pitch_rad)
        dy = dy * np.cos(roll_rad)

        # compute perpendicular offsets
        perp_dx = -dy
        perp_dy = dx

        # calculate left and right boundaries of the lane
        lane_width = waypoint.lane_width
        if noncross_left:
            noncross_boundary.append(
                [
                    waypoint.transform.location.x - lane_width / 2 * perp_dx,
                    waypoint.transform.location.y - lane_width / 2 * perp_dy,
                ]
            )
        if noncross_right:
            noncross_boundary.append(
                [
                    waypoint.transform.location.x + lane_width / 2 * perp_dx,
                    waypoint.transform.location.y + lane_width / 2 * perp_dy,
                ]
            )

    return noncross_boundary


def get_obstacle_2d_bbox_list(ego_pose, obstacle_predictions, waypoints):
    """Get obstacle 2D bounding box list"""
    if len(obstacle_predictions) == 0 or len(waypoints) == 0:
        return np.empty((0, 4))

    [ego_x, ego_y, _, _, _, ego_yaw] = ego_pose

    obstacle_list = []
    distances = pairwise_distances(waypoints, obstacle_predictions[:, :2]).min(0)

    for distance, prediction in zip(distances, obstacle_predictions):
        [x, y, _, confidence, label] = prediction
        angle = np.arctan2(y - ego_y, x - ego_x)
        diff_angle = np.arctan2(np.sin(angle - ego_yaw), np.cos(angle - ego_yaw))

        if (
            distance <= OBSTACLE_DIST_THRESHOLD
            and abs(diff_angle) <= OBSTACLE_ANGLE_THRESHOLD
            and confidence >= OBSTACLE_CONFIDENCE_THRESHOLD
            and LABELS[int(label)] not in OBSTACLE_IGNORE_LABELS
        ):
            # TODO: replace this with obstacle bbox
            obstacle_size = np.array(
                [
                    x
                    - OBSTACLE_RADIUS * np.cos(angle) / 2
                    - OBSTACLE_RADIUS_TANGENT * np.sin(angle) / 2,
                    y
                    - OBSTACLE_RADIUS * np.sin(angle) / 2
                    - OBSTACLE_RADIUS_TANGENT * np.cos(angle) / 2,
                    x
                    + OBSTACLE_RADIUS * np.cos(angle) / 2
                    + OBSTACLE_RADIUS_TANGENT * np.sin(angle) / 2,
                    y
                    + OBSTACLE_RADIUS * np.sin(angle) / 2
                    + OBSTACLE_RADIUS_TANGENT * np.cos(angle) / 2,
                    # x - OBSTACLE_RADIUS,
                    # y - OBSTACLE_RADIUS,
                    # x + OBSTACLE_RADIUS,
                    # y + OBSTACLE_RADIUS,
                ]
            )

            obstacle_list.append(obstacle_size)

    if len(obstacle_list) == 0:
        return np.empty((0, 4))

    return np.array(obstacle_list)


def get_lane_obstacle_2d_bbox_list(ego_pose, lane_obstacles):
    """Get lane 2D bounding box list with merged obstacle lines"""
    if len(lane_obstacles) == 0:
        return np.empty((0, 4))

    merged_obstacle_list = []
    i = 0
    while i < len(lane_obstacles):
        segment_start = lane_obstacles[i]
        segment_end = lane_obstacles[i]
        i += 1

        if i < len(lane_obstacles):
            segment_angle = np.arctan2(
                lane_obstacles[i][1] - segment_start[1],
                lane_obstacles[i][0] - segment_start[0],
            )

        while i < len(lane_obstacles):
            next_point = lane_obstacles[i]
            next_angle = np.arctan2(
                next_point[1] - segment_end[1], next_point[0] - segment_end[0]
            )

            if np.linalg.norm(next_point - segment_end) <= 1.5 and np.isclose(
                next_angle, segment_angle, atol=np.deg2rad(5)
            ):
                segment_end = next_point
                i += 1
            else:
                break

        if np.array_equal(segment_start, segment_end):
            bbox = np.array(
                [
                    segment_start[0] - OBSTACLE_LANE_RADIUS,
                    segment_start[1] - OBSTACLE_LANE_RADIUS,
                    segment_start[0] + OBSTACLE_LANE_RADIUS,
                    segment_start[1] + OBSTACLE_LANE_RADIUS,
                ]
            )
            merged_obstacle_list.append(bbox)
        else:
            segment_length = np.linalg.norm(segment_end - segment_start)
            middle_point = (segment_start + segment_end) / 2
            bbox = np.array(
                [
                    middle_point[0]
                    - segment_length / 2 * np.cos(segment_angle)
                    - OBSTACLE_LANE_RADIUS * np.sin(segment_angle),
                    middle_point[1]
                    - segment_length / 2 * np.sin(segment_angle)
                    - OBSTACLE_LANE_RADIUS * np.cos(segment_angle),
                    middle_point[0]
                    + segment_length / 2 * np.cos(segment_angle)
                    + OBSTACLE_LANE_RADIUS * np.sin(segment_angle),
                    middle_point[1]
                    + segment_length / 2 * np.sin(segment_angle)
                    + OBSTACLE_LANE_RADIUS * np.cos(segment_angle),
                ]
            )
            merged_obstacle_list.append(bbox)

    if len(merged_obstacle_list) == 0:
        return np.empty((0, 4))

    return np.array(merged_obstacle_list)


class VehicleState:
    """
    Vehicle State
    """

    def __init__(self):
        self.pose = None
        self.speed = None
        self.yaw_rate = None
        self.accel_body = None
        self.action = Action.NONE
        self.status = Status.IDLE

    def is_ready(self):
        if (
            self.pose is None
            or self.speed is None
            or self.yaw_rate is None
            or self.accel_body is None
        ):
            return False

        return True

    def set_pose(self, position):
        if len(position) == 0:
            return

        [x, y, z, qx, qy, qz, qw] = position
        [roll, pitch, yaw] = R.from_quat([qx, qy, qz, qw]).as_euler(
            "xyz", degrees=False
        )
        self.pose = np.array([x, y, z, roll, pitch, yaw])

    def set_velocity(self, speed_vector):
        if len(speed_vector) == 0:
            return

        self.speed = max(speed_vector[0], 0.01)

    def set_imu(self, imu_vector):
        if len(imu_vector) == 0:
            return

        self.yaw_rate = imu_vector[5]
        self.accel_body = imu_vector[:3].copy()

    def set_status(self, status: Status):
        self.status = status

    def set_action(self, action: Action):
        self.action = action

    def get_action(self):
        if not self.is_ready():
            return

        return self.action

    def get_status(self):
        if not self.is_ready():
            return

        return self.status

    def get_xy(self):
        if not self.is_ready():
            return

        return self.pose[:2]

    def get_xyz(self):
        if not self.is_ready():
            return

        return self.pose[:3]

    def get_yaw(self):
        if not self.is_ready():
            return

        return wrap_angle(self.pose[5])

    def get_speed(self):
        if not self.is_ready():
            return

        return self.speed

    def get_forward_vector(self):
        if not self.is_ready():
            return

        forward_vector = np.array(
            [
                np.cos(self.pose[5]) * np.cos(self.pose[4]),
                np.sin(self.pose[4]),
                np.sin(self.pose[5]) * np.cos(self.pose[4]),
            ]
        )
        return forward_vector / np.linalg.norm(forward_vector)

    def get_vel_vector(self, body_frame=False):
        if not self.is_ready():
            return

        if body_frame:
            return np.array([self.speed, 0, 0])
        else:
            # rot_matrix = R.from_euler(
            #     "xyz", [self.pose[3], self.pose[4], self.pose[5]], degrees=False
            # ).as_matrix()
            rot_matrix = R.from_euler(
                "xyz", [0, 0, self.pose[5]], degrees=False
            ).as_matrix()
            vel_vec = rot_matrix @ np.array([self.speed, 0, 0]).flatten()

            return vel_vec[:2]

    def get_pose(self):
        if not self.is_ready():
            return

        return self.pose

    def get_accel_vector(self, body_frame=True):
        if not self.is_ready():
            return

        if body_frame:
            return self.accel_body[:2]
        else:
            # rot_matrix = R.from_euler(
            #     "xyz", [self.pose[3], self.pose[4], self.pose[5]], degrees=False
            # ).as_matrix()
            rot_matrix = R.from_euler(
                "xyz", [0, 0, self.pose[5]], degrees=False
            ).as_matrix()
            accel_vec = rot_matrix @ self.accel_body.flatten()

            return accel_vec[:2]

    def predict_next_pose(self, dt=VEHICLE_DT):
        if not self.is_ready():
            return

        next_pose = self.pose.copy()
        next_pose[0] = self.pose[0] + np.cos(self.pose[5]) * self.speed * dt
        next_pose[1] = self.pose[1] + np.sin(self.pose[5]) * self.speed * dt
        next_pose[5] = wrap_angle(self.pose[5] + self.yaw_rate * dt)

        return next_pose


class Operator:
    """Send the frenet optimal trajectory `waypoints` given some initial conditions in
    cartesian space.
    Args:
        initial_conditions (dict): dict containing the following items
            ps (float): previous longitudinal position
            target_speed (float): target speed [m/s]
            pos (np.ndarray([float, float])): initial position in global coord
            vel (np.ndarray([float, float])): initial velocity [m/s]
            accel (np.ndarray([float, float])): initial acceleration in global coord [m/s^2]
            wp (np.ndarray([float, float])): list of global waypoints
            obs (np.ndarray([float, float, float, float])): list of obstacles
                as: [lower left x, lower left y, upper right x, upper right y]
        hyperparameters (dict): a dict of optional hyperparameters
            max_speed (float): maximum speed [m/s]
            max_accel (float): maximum acceleration [m/s^2]
            max_curvature (float): maximum curvature [1/m]
            max_road_width_l (float): maximum road width to the left [m]
            max_road_width_r (float): maximum road width to the right [m]
            d_road_w (float): road width sampling discretization [m]
            dt (float): time sampling discretization [s]
            maxt (float): max prediction horizon [s]
            mint (float): min prediction horizon [s]
            d_t_s (float): target speed sampling discretization [m/s]
            n_s_sample (float): sampling number of target speed
            obstacle_clearance (float): obstacle radius [m]
            kd (float): positional deviation cost
            kv (float): velocity cost
            ka (float): acceleration cost
            kj (float): jerk cost
            kt (float): time cost
            ko (float): dist to obstacle cost
            klat (float): lateral cost
            klon (float): longitudinal cost
    Returns:
        result_x (np.ndarray(float)): x positions of fot, if it exists
        result_y (np.ndarray(float)): y positions of fot, if it exists
        speeds (np.ndarray(float)): speeds of fot, if it exists
        ix (np.ndarray(float)): spline x of fot, if it exists
        iy (np.ndarray(float)): spline y of fot, if it exists
        iyaw (np.ndarray(float)): spline yaws of fot, if it exists
        d (np.ndarray(float)): lateral offset of fot, if it exists
        s (np.ndarray(float)): longitudinal offset of fot, if it exists
        speeds_x (np.ndarray(float)): x speeds of fot, if it exists
        speeds_y (np.ndarray(float)): y speeds of fot, if it exists
        params (dict): next frenet coordinates, if they exist
        costs (dict): costs of best frenet path, if it exists
        success (bool): whether a fot was found or not
    """

    def __init__(self):
        self.vehicle_state = VehicleState()
        self.obstacles = np.array([])
        self.all_obstacles = np.array([])
        self.failure_counter = 0
        self.last_local_trajectory = None

        if LOCAL_PLANNER_USE_LANE_OBSTACLES:
            self.hd_map = None
            self.all_waypoints = []
            self.noncross_boundary = np.array([])

        # FOT hyperparameters
        self.hyperparameters = {
            "max_speed": LOCAL_PLANNER_MAX_SPEED,
            "max_accel": LOCAL_PLANNER_MAX_ACCEL,
            "max_curvature": 5.0,
            "max_road_width_l": 0.5,
            "max_road_width_r": 0.5,
            "d_road_w": 0.1,
            "dt": 0.1,
            "maxt": 5.0,
            "mint": 2.0,
            "d_t_s": 0.2,
            "n_s_sample": 20.0,
            "obstacle_clearance": 0.25,
            "kd": 10.0,
            "kv": 0.1,
            "ka": 0.1,
            "kj": 1.0,
            "kt": 0.1,
            "ko": 5.0,
            "klat": 1.0,
            "klon": 1.0,
            "num_threads": 6,  # set 0 to avoid using threaded algorithm
        }

        ## Debug
        if DEBUG:
            self.fig, self.axis = plt.subplots(1, 1)
            self.fig.canvas.set_window_title("FOT Monitor")
            self.fig.canvas.manager.window.wm_geometry("+100+800")
            (self.ego_handle,) = self.axis.plot(
                [], [], "b^", markersize=10, label="ego"
            )
            (self.fot_xy_handle,) = self.axis.plot(
                [], [], "o", markersize=5, color="C1", alpha=0.6, label="fot_xy"
            )
            (self.global_waypoints_handle,) = self.axis.plot(
                [], [], "s", markersize=5, color="g", alpha=0.6, label="global_wp"
            )
            self.axis.legend()
            self.axis.set_aspect(1)

    def on_event(
        self,
        dora_event: dict,
        send_output: Callable[[str, bytes], None],
    ) -> DoraStatus:
        if dora_event["type"] == "INPUT":
            return self.on_input(dora_event, send_output)
        return DoraStatus.CONTINUE

    def on_input(
        self,
        dora_input: dict,
        send_output: Callable[[str, bytes], None],
    ):
        if dora_input["id"] == "position":
            position = np.frombuffer(dora_input["data"], np.float32)
            self.vehicle_state.set_pose(position)
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "speed":
            speed = np.frombuffer(dora_input["data"], np.float32)
            self.vehicle_state.set_velocity(speed)
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "imu":
            imu = np.frombuffer(dora_input["data"], np.float32)
            self.vehicle_state.set_imu(imu)
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "opendrive":
            if LOCAL_PLANNER_USE_LANE_OBSTACLES and self.hd_map is None:
                opendrive = dora_input["data"].decode()
                self.hd_map = carla.Map("map", opendrive)
                self.all_waypoints = self.hd_map.generate_waypoints(distance=1.0)
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "obstacles":
            obstacles = np.frombuffer(dora_input["data"], np.float32).reshape((-1, 5))
            self.obstacles = obstacles
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "decision_making":
            decision = Decision.from_bytes(dora_input["data"])
            self.vehicle_state.set_action(decision.action)
            self.vehicle_state.set_status(decision.status)
            return DoraStatus.CONTINUE

        elif dora_input["id"] == "global_waypoints":
            waypoints = np.frombuffer(dora_input["data"], np.float32).reshape(-1, 4)

            # If no localization or speed measurement
            # This should only happen at initialization
            if not self.vehicle_state.is_ready():
                return DoraStatus.CONTINUE

            # If global waypoints are empty
            if len(waypoints) == 0:
                local_trajectory = self.reset_trajectory()
                self.vehicle_state.set_action(Action.STOP)

                send_output(
                    "waypoints",
                    local_trajectory[:, :3].tobytes(),
                    {},
                )

                send_output(
                    "trajectory",
                    local_trajectory.tobytes(),
                    {},
                )
                return DoraStatus.CONTINUE

            # Get global waypoints
            global_waypoints_xyz = waypoints[:, :3][
                : min(LOCAL_PLANNER_NUM_GLOBAL_WAYPOINTS_AHEAD, len(waypoints))
            ]
            global_waypoints = global_waypoints_xyz[:, :2]

            # Get target speed from the last global waypoint
            target_speed = waypoints[-1, -1]
            if target_speed>0:
                target_speed+=5

            # Get static obstacles
            static_obstacles = get_obstacle_2d_bbox_list(
                self.vehicle_state.get_pose(), self.obstacles, global_waypoints
            )

            # Decide whether use last trajectory
            use_last_trajectory = False
            if (
                self.last_local_trajectory is not None
                and len(self.last_local_trajectory) > 2
                and self.vehicle_state.get_action() == Action.NONE
            ):
                closest_point_index = self.get_closest_point_index(
                    self.vehicle_state.get_xy(), self.last_local_trajectory[:, :2]
                )

                # Point does not exist
                if closest_point_index is None:
                    use_last_trajectory = False

                # Only one point left
                elif closest_point_index == len(self.last_local_trajectory) - 1:
                    use_last_trajectory = False

                # Tracking error is larger than 0.5 m
                elif (
                    np.linalg.norm(
                        self.vehicle_state.get_xy()
                        - self.last_local_trajectory[closest_point_index, :2],
                    )
                    > 0.5
                ):
                    use_last_trajectory = False

                # Time left is smaller than 2.0 second
                elif len(self.last_local_trajectory) - closest_point_index < int(
                    2.0 / self.hyperparameters["dt"]
                ):
                    use_last_trajectory = False

                # Distance left is smaller than 2.0 m
                elif (
                    np.linalg.norm(
                        self.vehicle_state.get_xy() - self.last_local_trajectory[-1, :2]
                    )
                    < 2.0
                ):
                    use_last_trajectory = False

                else:
                    use_last_trajectory = True

            # If using last trajectory
            if use_last_trajectory:
                self.last_local_trajectory = self.last_local_trajectory[
                    closest_point_index:, :
                ]
                local_trajectory = self.last_local_trajectory.copy()
            else:
                # replanning flag
                replan_use_actual_state = False

                # When last trajectory does not exist
                if self.last_local_trajectory is None:
                    replan_use_actual_state = True

                # When the vehicle just starts
                elif abs(self.vehicle_state.get_speed()) < 0.5:
                    replan_use_actual_state = True

                # When last trajectory is too short
                elif (
                    int(np.ceil(VEHICLE_DT / self.hyperparameters["dt"]))
                    > len(self.last_local_trajectory) - 1
                ):
                    replan_use_actual_state = True

                # When the error is too large
                elif (
                    np.linalg.norm(
                        self.vehicle_state.get_xy()
                        - self.last_local_trajectory[
                            int(np.ceil(VEHICLE_DT / self.hyperparameters["dt"])), :2
                        ]
                    )
                    > 0.1
                ):
                    replan_use_actual_state = True

                # If replan using actual state
                replan_use_actual_state = True
                if replan_use_actual_state:
                    if abs(self.vehicle_state.get_speed()) < 0.5:
                        init_pos = self.vehicle_state.get_xy()
                        init_vel = self.vehicle_state.get_vel_vector()
                    else:
                        next_pose = self.vehicle_state.predict_next_pose()
                        init_pos = next_pose[:2]
                        init_vel = np.array(
                            [
                                self.vehicle_state.get_speed() * np.cos(next_pose[-1]),
                                self.vehicle_state.get_speed() * np.sin(next_pose[-1]),
                            ]
                        )
                else:
                    next_index = int(np.ceil(VEHICLE_DT / self.hyperparameters["dt"]))
                    init_pos = self.last_local_trajectory[next_index, :2]
                    init_vel = np.array(
                        [
                            self.last_local_trajectory[next_index, 2]
                            * np.cos(self.last_local_trajectory[next_index, 3]),
                            self.last_local_trajectory[next_index, 2]
                            * np.sin(self.last_local_trajectory[next_index, 3]),
                        ]
                    )

                init_accel = self.vehicle_state.get_accel_vector(body_frame=False)

                # Get lane obstacles
                if LOCAL_PLANNER_USE_LANE_OBSTACLES and (self.hd_map is not None):
                    road_lane_id_list = []
                    for global_waypoint_xyz in global_waypoints_xyz:
                        [_x, _y, _z] = global_waypoint_xyz
                        waypoint = self.hd_map.get_waypoint(
                            carla.Location(float(_x), float(_y), float(_z)),
                            project_to_road=True,
                            lane_type=carla.LaneType.Driving,
                        )

                        if (
                            waypoint.road_id,
                            waypoint.lane_id,
                        ) not in road_lane_id_list:
                            road_lane_id_list.append(
                                (waypoint.road_id, waypoint.lane_id)
                            )

                    # get lane waypoints
                    lane_waypoints = []
                    for wp in self.all_waypoints:
                        if wp.is_junction:
                            continue
                        if (wp.road_id, wp.lane_id) not in road_lane_id_list:
                            continue
                        wp_xyz = np.array(
                            [
                                wp.transform.location.x,
                                wp.transform.location.y,
                                wp.transform.location.z,
                            ]
                        )

                        # only select those close to the global waypoints
                        min_distance = pairwise_distances(
                            global_waypoints,
                            wp_xyz[:2].reshape(1, -1),
                        ).min(0)
                        if min_distance > 1.0:
                            continue

                        # only select those near the vehicle:
                        distance_to_ego = np.linalg.norm(
                            wp_xyz[:2] - self.vehicle_state.get_xy()
                        )
                        if distance_to_ego > 10.0:
                            continue

                        # only select those in front of the vehicle:
                        forward_vec = self.vehicle_state.get_forward_vector()
                        ray = wp_xyz - self.vehicle_state.get_xyz()
                        if forward_vec.dot(ray) < 0:
                            continue

                        lane_waypoints.append(wp)

                    # get lane boundaries
                    if lane_waypoints is not None:
                        noncross_boundary = get_noncross_boundary(lane_waypoints)
                        self.noncross_boundary = np.asarray(noncross_boundary)

                    # get lane obstacles
                    lane_obstacles = np.empty((0, 2))
                    if len(self.noncross_boundary) > 0:
                        lane_obstacles = np.concatenate(
                            (lane_obstacles, self.noncross_boundary)
                        )

                    lane_obstacles = get_lane_obstacle_2d_bbox_list(
                        self.vehicle_state.get_pose(), lane_obstacles
                    )

                    self.all_obstacles = np.concatenate(
                        [static_obstacles, lane_obstacles]
                    )

                else:
                    self.all_obstacles = static_obstacles

                # FOT initial conditions
                initial_conditions = {
                    "ps": 1,
                    "target_speed": np.clip(target_speed, 0, LOCAL_PLANNER_MAX_SPEED),
                    "pos": init_pos,
                    "vel": init_vel,
                    "accel": init_accel,
                    "wp": global_waypoints,
                    "obs": self.all_obstacles,
                }

                # FOT run
                (
                    result_x,
                    result_y,
                    speeds,
                    ix,
                    iy,
                    iyaw,
                    d,
                    s,
                    speeds_x,
                    speeds_y,
                    misc,
                    costs,
                    success,
                ) = fot_wrapper.run_fot(initial_conditions, self.hyperparameters)

                if DEBUG:
                    output = {
                        "initial_conditions": initial_conditions,
                        "use_last_trajectory": use_last_trajectory,
                        "last_trajectory": self.last_local_trajectory,
                        "replan_use_actual_state": replan_use_actual_state,
                        "vehicle_pose": self.vehicle_state.get_pose(),
                        "vehicle_speed": self.vehicle_state.get_speed(),
                        "vehicle_yawrate": self.vehicle_state.yaw_rate,
                        # "action": self.vehicle_state.get_action(),
                        "success": success,
                    }
                    with open(filename, "ab") as file:
                        pickle.dump(output, file)

                if not success:
                    self.failure_counter += 1
                    if (
                        self.failure_counter < 3
                        and self.last_local_trajectory is not None
                    ):
                        local_trajectory = self.last_local_trajectory.copy()
                        if DEBUG:
                            print(
                                "[FOT] Compute local waypoints failed! Use last computed result."
                            )
                    else:
                        # try to follow the point on the global waypoints instead
                        local_trajectory = np.zeros((1, 7), np.float32)

                        closest_point_index = self.get_closest_point_index(
                            self.vehicle_state.get_xy(), global_waypoints
                        )
                        lookahead_point_index = min(
                            closest_point_index + 4, len(global_waypoints) - 1
                        )
                        current_yaw = self.vehicle_state.get_yaw()
                        local_trajectory[0, 0] = global_waypoints[
                            lookahead_point_index, 0
                        ]
                        local_trajectory[0, 1] = global_waypoints[
                            lookahead_point_index, 1
                        ]
                        local_trajectory[0, 2] = target_speed
                        local_trajectory[0, 3] = current_yaw
                        self.last_local_trajectory = local_trajectory.copy()

                        self.failure_counter = 0

                        if DEBUG:
                            print(
                                "[FOT] Compute local waypoints failed! Follow global waypoints instead."
                            )

                else:
                    local_trajectory = np.zeros((len(s), 7))
                    local_trajectory[:, 0] = result_x
                    local_trajectory[:, 1] = result_y
                    local_trajectory[:, 2] = speeds
                    local_trajectory[:, 3] = iyaw
                    local_trajectory[:, 5] = s
                    local_trajectory[:, 6] = np.array(
                        [i * VEHICLE_DT for i in range(len(s))]
                    )

                    # calculate curvature
                    if len(local_trajectory) == 1:
                        local_trajectory[:, 4] = np.array([0.0])
                    else:
                        ds = np.hypot(np.diff(result_x), np.diff(result_y))
                        local_trajectory[:-1, 4] = wrap_angle(np.diff(iyaw)) / ds
                        local_trajectory[-1, 4] = local_trajectory[-2, 4]

                    local_trajectory = np.around(local_trajectory, 5).astype(np.float32)
                    self.last_local_trajectory = local_trajectory.copy()

                    # reset counter
                    self.failure_counter = 0

            ## Debug
            if DEBUG:
                self.ego_handle.set_xdata(self.vehicle_state.get_xy()[0])
                self.ego_handle.set_ydata(self.vehicle_state.get_xy()[1])

                self.fot_xy_handle.set_xdata(local_trajectory[:, 0])
                self.fot_xy_handle.set_ydata(local_trajectory[:, 1])

                self.global_waypoints_handle.set_xdata(global_waypoints[:, 0])
                self.global_waypoints_handle.set_ydata(global_waypoints[:, 1])

                for patch in self.axis.patches:
                    patch.remove()

                for obstacle in self.all_obstacles:
                    [x_min, y_min, x_max, y_max] = obstacle
                    vertices = np.array(
                        [
                            [x_min, y_min],
                            [x_max, y_min],
                            [x_max, y_max],
                            [x_min, y_max],
                        ]
                    )
                    self.axis.add_patch(
                        plt.Polygon(
                            vertices, facecolor="k", edgecolor="none", alpha=0.85
                        )
                    )

                self.axis.set_title(
                    f"current speed: {self.vehicle_state.get_speed():.1f}, target speed: {target_speed:.1f}, action: {self.vehicle_state.get_action()}, # of obs: {len(self.all_obstacles)}",
                    fontsize=11,
                )
                self.axis.set_xlim(
                    [
                        self.vehicle_state.get_xy()[0] - 20,
                        self.vehicle_state.get_xy()[0] + 20,
                    ]
                )
                self.axis.set_ylim(
                    [
                        self.vehicle_state.get_xy()[1] - 20,
                        self.vehicle_state.get_xy()[1] + 20,
                    ]
                )
                self.axis.invert_yaxis()
                plt.pause(1e-5)

            # If emergency stop
            if self.vehicle_state.get_action() == Action.STOP:
                local_trajectory = self.reset_trajectory()

            send_output(
                "waypoints",
                local_trajectory[:, :3].tobytes(),  # x, y, vel
                {},
            )

            send_output(
                "trajectory",
                local_trajectory.tobytes(),  # x, y, vel, yaw, k, s, t
                {},
            )
            return DoraStatus.CONTINUE

        return DoraStatus.CONTINUE

    def reset_trajectory(self):
        # trajectory = [x, y, vel, yaw, k, s, t]
        local_trajectory = np.empty((0, 7), np.float32)
        self.last_local_trajectory = None

        return local_trajectory

    @staticmethod
    def get_closest_point_index(position, waypoints):
        if waypoints is None or len(waypoints) < 1:
            return

        distances = pairwise_distances(
            waypoints, np.asarray(position[:2]).reshape(-1, 2)
        ).T[0]

        closest_point_index = np.argmin(distances)
        return closest_point_index
